#pragma once
/*
 *	Copyright © 2008 University of Houston
 *	All rights reserved.
 */

#include <time.h>
#include <pcap.h>
#include <vector>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <ctype.h>
#include <errno.h>
#include <string>
#include "sniffex.h"
#include "util.h"
#include "analyzer.h"
using namespace std;

class session;
class sessionAnalyzer;
class timedHeader;

typedef vector<session*> 	session_list;
typedef vector<timedHeader>	header_list;

/*
 * Encapsulates a TCP packet header, incorporating timestamp information
 * regarding the time that the packet was either sent or received.
 */
class timedHeader
{
public:
	timedHeader();
	~timedHeader();

	/*
	 * Formats raw packet data into an intelligble structure, usable
	 * by the timedHeader class.
	 * @param packet Pointer to the packet data.
	 * @return False if an error occurs, otherwise true.
	 */
	bool formatPacket(const u_char* packet);

	/*
	 * Returns an in_addr structure that represents the IP address of
	 * the originating (sending) machine.
	 * @return An in_addr structure representing the IP address.
	 */
	in_addr		getSourceIP();

	/*
	 * Returns an in_addr structure that represents the IP address of
	 * the receiving machine.
	 * @return An in_addr structure representing the IP address.
	 */
	in_addr		getDestIP();

	/*
	 * Returns the port over which the connection has been made on the
	 * source machine.
	 * @return The port.
	 */
	int			getSourcePort();

	/*
	 * Returns the port over which the connection has been made on the
	 * destination machine.
	 * @return The port.
	 */
	int			getDestPort();

	/*
	 * Returns the size (in bytes) of the packet payload.
	 * @return The payload size in bytes.
	 */
	int			getPayloadSize();

	/*
	 * Returns the timestamp of the packet as it was received or sent
	 * by the local machine.  The timestamp is represented in fractions
	 * of a second.  A timestamp value of 1 is equal to
	 * 1/CLOCKS_PER_SEC seconds.  The timestamp value starts when the
	 * program starts.
	 * @return The timestamp.
	 */
	timeval	getTimestamp();

	/*
	 * @return Returns the time (in secons) since the program started when the
	 * packet was received.
	 * @return The time, in seconds.
	 */
	unsigned long	getTime();

	/*
	 * Returns a pointer to the internal sniff_tcp object used to
	 * determine information about the packet.
	 * @return A pointer to a sniff_tcp object.
	 */
	sniff_tcp*	getTCP();

	/*
	 * Returns a pointer to the internal sniff_ip object used to
	 * determine information about the packet.
	 * @return A pointer to a sniff_ip object.
	 */
	sniff_ip*	getIP();

	bool	isFinAck();
	bool	isAck();
	bool	isFin();
	bool	isPush();

	timedHeader	getAck(timedHeader pshPacket);
	timedHeader	getPsh(timedHeader ackPacket);

	tcp_seq	getSequence();
	tcp_seq getAck();

	// Operator overloads
	timedHeader& operator=(const u_char*);
	timedHeader& operator=(const timedHeader);

	bool 	seqMatches(timedHeader other);

	// The session, packetLogger classes benefit from having direct access.
	friend class session;
	friend class packetLogger;
private:
	void setSrcIP(const char* srcIPBuffer);
	void setDstIP(const char* dstIPBuffer);
	void setSrcPort(const u_int srcPort);
	void setDstPort(const u_int dstPort);
	void setPayloadSize(const int payloadSize);
	void setFlags(const u_char flags);

	pcap_pkthdr		packetHeader;
	sniff_ip		ipHeader;
	sniff_tcp		tcpHeader;
	sniff_ethernet	ethHeader;
	int				payloadSize;

	timeval			timestamp;
};




/*
 * Encapsulates a single connection (e.g. a SSH session), storing all of
 * the pertinent packets, as well as the IP information.
 */
class session
{
public:
	session();
	session(timedHeader packet);
	~session();

	/*
	 * Initializes the variables used to detect session closing.
	 * Does not need to be called explicitly.
	 */
	void		initialize();

	/*
	 * Adds a packet to our internal list of packets.
	 * @param packet The packet to add
	 * @return No return value.
	 */
	void		addPacket(timedHeader packet);
	void		loadPacket(timedHeader th);

	/*
	 * Returns the number of packets that have been stored/recorded.
	 * @return The number of packets that have been stored/recorded.
	 */
	int 		getNumPackets() const;

	/*
	 * Get the vector of packets.
	 * @return headerList
	 */
	vector<timedHeader>* getHeaderList();

	/*
	 * Returns a packet based on its packet ID #.  Packet ID #'s are
	 * sequential, starting with zero.  The first packet we receive will
	 * have an ID # of 0, and the most recent packet will have an ID # of
	 * getNumPackets()-1.
	 * @param offset The packet ID #
	 * @return A copy of the packet header structure.
	 */
	timedHeader	getPacket(int offset) const;

	/*
	 * Sets the Analyzer that this session will use.
	 * @param pAnalyzer Pointer to an Analyzer object to be used.
	 */
	void		setAnalyzer(sessionAnalyzer* pAnalyzer);

	/*
	 * @return Returns a copy of the most recent packet received.
	 */
	timedHeader	getLatestPacket() const;

	/*
	 * @return Returns true if we have successfully observed the full 3-way
	 * handeshake for connection closing.
	 * A->B FIN/ACK
	 * B->A FIN/ACK
	 * A->B ACK
	 */
	bool 		isConnectionClosed() const;

	/*
	 * @return Returns the session ID of this session.  Used for log tracking.
	 */
	int 		getSessionID() const;

	/*
	 * Set the session ID for this session.  Used for log tracking.
	 * @param newID The ID to set.
	 */
	void 		setSessionID(int newID);


	// Methods to retrieve the IP address of host A and host B, and the ports
	in_addr		getIPA() const;
	in_addr		getIPB() const;

	// Get ports...
	u_short		getPortA() const;
	u_short		getPortB() const;

	// Methods to check and see if a packet belongs to this session.
	bool		belongs(timedHeader param) const;

	timedHeader	operator[](const unsigned int offset) const;

	friend class packetLogger;
private:
	// Track connection state when closing...
	bool		finack1, finack2, final_ack;
	int			finack1_index, finack2_index, final_ack_index;

	header_list		headerList;

	int			sessionID;
	//filter*		pFilter;
	sessionAnalyzer*	pAnalyzer;
};

